# use focal loss for binary classfication by xgb

from imxgboost.imbalance_xgb import imbalance_xgboost as imb_xgb
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix, recall_score, precision_recall_curve
import numpy as np
import pandas as pd


def read_data():
    data = pd.read_csv('D:/lung_cancer/data/data.csv')
    features = []
    labels = []
    for i in range(len(data)):
        if data['cancer_type'][i]-1 == 0 or data['cancer_type'][i]-1 == 1:
            one_feature = [data['z'][i], data['x'][i], data['y'][i], data['r'][i],
                           data['patientWeight'][i], data['patientSex'][i], data['patientAge'][i],
                           data['part_suvmax'][i], data['suv_avg'][i], data['suv_std'][i]]
            features.append(one_feature)
            labels.append(data['cancer_type'][i]-1)
    print(features)
    print(len(features))
    print(labels)
    print(len(labels))

    # 只用第一类和第二类数据

    X_train, X_test, y_train, y_test = train_test_split(np.array(features, dtype=np.float),
                                                        np.array(labels, dtype=np.int),
                                                        test_size=0.2, random_state=1234565)

    print('Xtrain type: ', type(X_train))
    print('y_train type: ', type(y_train))
    return X_train, X_test, y_train, y_test

def train():
    X_train, X_test, y_train, y_test = read_data()
    xgboster_focal = imb_xgb(special_objective='focal')
    xgboster_weight = imb_xgb(special_objective='weighted')
    CV_focal_booster = GridSearchCV(xgboster_focal, {"focal_gamma":[1.0, 1.5, 2.0, 2.5, 3.0]})
    CV_weight_booster = GridSearchCV(xgboster_weight, {"imbalance_alpha":[1.5, 2.0, 2.5, 3.0, 4.0]})
    CV_focal_booster.fit(X_train, y_train)
    CV_weight_booster.fit(X_train, y_train)
    opt_focal_booster = CV_focal_booster.best_estimator_
    opt_weight_booster = CV_weight_booster.best_estimator_

    raw_output = opt_focal_booster.predict(X_test, y=None)
    sigmoid_output = opt_focal_booster.predict_sigmoid(X_test, y=None)
    class_output = opt_focal_booster.predict_determine(X_test, y=None)
    prob_output = opt_focal_booster.predict_two_class(X_test, y=None)

    print(class_output)
    # 计算评价指标：混淆矩阵+准确率+查准+查全+F1-score
    test_accuracy = accuracy_score(y_test, class_output)
    print("test accuracy: %.2f %%" % (test_accuracy * 100))

    target_names = ['腺癌', '鳞癌']
    result_statis = classification_report(y_test, class_output, target_names=target_names)
    print(result_statis)

    confusion = confusion_matrix(y_test, class_output)
    print(confusion)

    la = precision_recall_curve(y_test, class_output)
    print(la)



if __name__ == '__main__':
    train()
